﻿#include "VulkanDirectionalLights.h"
#include "Render/Widget/VulkanViewportWidget.h"
#include "Render/Scene/VulkanScene.h"
#include "Render/Scene/DefferRenderAttachments.h"
#include "Render/Texture/GPURenderTarget.h"
#include "Shaders/CppMacros.h"
#include "Render/Mesh/VulkanMesh.h"

void CVulkanDirectionalLights::PreProcess(
    const TVector<SPreRenderDirectionalLight>& InLightData
    , CVulkanViewportWidget * InViewport
    , CVulkanScene * InScene)
{
    LightDatas.Clear();

    for(auto& InLightData : InLightData)
    {
        //检查是否是黑色光源
        if(InLightData.Color.R==0.f && InLightData.Color.G==0.f && InLightData.Color.B==0.f)
        {
            continue;
        }

        LightDatas.Add(InLightData);

    }

    //绑定数据
    DirectionalLights.Resize(LightDatas.Num());
    for(int32_t LightIndex=0;LightIndex<LightDatas.Num();++LightIndex)
    {
        DirectionalLights[LightIndex].SetData(LightDatas[LightIndex]);
        DirectionalLights[LightIndex].LightIndex=LightIndex;
    }

    //统计描述符集数量
    SVulkanDescriptorPoolSize PoolSize;
    PoolSize.SetNum= LightDatas.Num();
    PoolSize+=CVulkanContext::Get()->GetSceneGlobal()->DirectionalLightDefferShadingDescriptorSetLayout.DescriptorNums * LightDatas.Num();
    CVulkanContext::Get()->GetAutoResetDescriptorPool()->GenerallPool.PoolSize+=PoolSize;

}

void CVulkanDirectionalLights::RecordShadowMapCommand(
    vk::CommandBuffer InCmdBuffers
    , CVulkanScene *InScene
    , const SFrustumPoints & InFrustumPoints
    ,float Near
    ,float Far
    )
{
    for(auto& DirectionalLight : DirectionalLights)
    {
        DirectionalLight.RecordShadowMapCommand(
            InCmdBuffers
            ,InScene
            ,InFrustumPoints
            ,Near
            ,Far
            );
    }
}

void CVulkanDirectionalLights::RecordDefferShadingCommand(
    vk::CommandBuffer InCmdBuffers
    , CVulkanScene *InScene
    , const SSceneParameters &InParameters
    , CDefferRenderAttachments *Attachments)
{
    InCmdBuffers.bindPipeline(vk::PipelineBindPoint::eGraphics, CVulkanContext::Get()->GetSceneGlobal()->DirectionalLightDefferShadingPipeline);
    
    for(auto& DirectionalLight : DirectionalLights)
    {
        DirectionalLight.RecordDefferShadingCommand(InCmdBuffers,InScene,InParameters,Attachments);
    }
}

int64_t CVulkanDirectionalLights::GetDirectionalLightNum() const
{
    return DirectionalLights.Num();
}

bool CVulkanDirectionalLights::IsInLightRange(int64_t InLightIndex, const SVec3 &InSpherePosition, float InSphereRadius)
{
    //TODO 实现这个函数
    return true;
}

CVulkanDirectionalLight::CVulkanDirectionalLight()
{
    DefferShadingUniformBuffer=std::make_shared<TVulkanUniformBuffer<SDirectionalLightDefferShadingFragmentUniformBuffer>>();
}

CVulkanDirectionalLight::~CVulkanDirectionalLight()
{
    //清除帧缓冲区
    for(auto& FrameBuffer: CascadeShadowMapFrameBuffers)
    {
        CVulkanContext::Get()->GetDevice().destroyFramebuffer(FrameBuffer);
    }
}

void CVulkanDirectionalLight::SetData(const SPreRenderDirectionalLight &InLightData)
{
    LightData=InLightData;
}

void CVulkanDirectionalLight::RecordShadowMapCommand(
    vk::CommandBuffer InCmdBuffers
    , CVulkanScene *InScene
    ,const SFrustumPoints & InFrustumPoints,float Near,float Far
    )
{
    int ShadowMapNum= CVulkanContext::Get()->GetSceneGlobal()->GetCascadeNum();
    assert(ShadowMapNum>0);


    CascadeMaxDepths=LightData.CascadeMaxDepths;

    //每一级近平面比值
    TVector<float> CascadeNearDepthPercents;
    CascadeNearDepthPercents.Resize(ShadowMapNum);
    CascadeNearDepthPercents[0]= 0;
    for (uint8_t i = 1; i < ShadowMapNum; i++)
    {
        CascadeNearDepthPercents[i] =CascadeMaxDepths[i-1]/(Far-Near);
    }

    //每一级远平面比值
    TVector<float> CascadeFarDepthPercents;
    CascadeFarDepthPercents.Resize(ShadowMapNum);
    for (uint8_t i = 0; i < ShadowMapNum; i++)
    {
        CascadeFarDepthPercents[i] =CascadeMaxDepths[i]/(Far-Near);
    }


    //沿着相机的前向方向，分割视锥体
    TVector<SFrustumPoints> Frustums;
    Frustums.Resize(ShadowMapNum);

    for (uint8_t i = 0; i < ShadowMapNum; i++)
    {
        Frustums[i].LeftTopNear = InFrustumPoints.LeftTopNear + (InFrustumPoints.LeftTopFar - InFrustumPoints.LeftTopNear) *CascadeNearDepthPercents[i];

        Frustums[i].RightTopNear = InFrustumPoints.RightTopNear + (InFrustumPoints.RightTopFar - InFrustumPoints.RightTopNear) *  CascadeNearDepthPercents[i];;
        Frustums[i].LeftBottomNear = InFrustumPoints.LeftBottomNear + (InFrustumPoints.LeftBottomFar - InFrustumPoints.LeftBottomNear) *  CascadeNearDepthPercents[i];;
        Frustums[i].RightBottomNear = InFrustumPoints.RightBottomNear + (InFrustumPoints.RightBottomFar - InFrustumPoints.RightBottomNear) *  CascadeNearDepthPercents[i];

        Frustums[i].LeftTopFar = InFrustumPoints.LeftTopNear + (InFrustumPoints.LeftTopFar - InFrustumPoints.LeftTopNear) * CascadeFarDepthPercents[i];
        Frustums[i].RightTopFar = InFrustumPoints.RightTopNear + (InFrustumPoints.RightTopFar - InFrustumPoints.RightTopNear) *  CascadeFarDepthPercents[i];
        Frustums[i].LeftBottomFar = InFrustumPoints.LeftBottomNear + (InFrustumPoints.LeftBottomFar - InFrustumPoints.LeftBottomNear) * CascadeFarDepthPercents[i];
        Frustums[i].RightBottomFar = InFrustumPoints.RightBottomNear + (InFrustumPoints.RightBottomFar - InFrustumPoints.RightBottomNear) *  CascadeFarDepthPercents[i];
    }

    //project frustums into light space
    ProjectViewMatrixs.Resize(ShadowMapNum);

    SVec3 XAxis = LightData.Rotation.GetRightDirection();
    SVec3 YAxis = -LightData.Rotation.GetUpDirection();
    SVec3 ZAxis = LightData.Rotation.GetForwardDirection();

    //计算出每一级的投影矩阵
    for (uint8_t i = 0; i < ShadowMapNum; i++)
    {
        //直射光坐标系：
        //X轴是光源的右方向
        //Y轴是光源的上方向
        //Z轴是光源的前方向
        //原点是视锥体的近左上角(视锥体的其他点也可以)

        //视椎体的8个顶点在直射光坐标系下的坐标
        SVec3 FrustumPointInLightCoordinate[8];
        
        SVec3 Orignal = Frustums[i].LeftTopNear;

        //left top near
        FrustumPointInLightCoordinate[0] =  SVec3(0,0,0);

        //right top near
        {
            SVec3 OffsetToOrignal=Frustums[i].RightTopNear-Orignal;
            FrustumPointInLightCoordinate[1].X=OffsetToOrignal.DotProduct(XAxis);
            FrustumPointInLightCoordinate[1].Y=OffsetToOrignal.DotProduct(YAxis);
            FrustumPointInLightCoordinate[1].Z=OffsetToOrignal.DotProduct(ZAxis);
        }   

        //left bottom near
        {
            SVec3 OffsetToOrignal=Frustums[i].LeftBottomNear-Orignal;
            FrustumPointInLightCoordinate[2].X=OffsetToOrignal.DotProduct(XAxis);
            FrustumPointInLightCoordinate[2].Y=OffsetToOrignal.DotProduct(YAxis);
            FrustumPointInLightCoordinate[2].Z=OffsetToOrignal.DotProduct(ZAxis);
        }

        //right bottom near
        {
            SVec3 OffsetToOrignal=Frustums[i].RightBottomNear-Orignal;
            FrustumPointInLightCoordinate[3].X=OffsetToOrignal.DotProduct(XAxis);
            FrustumPointInLightCoordinate[3].Y=OffsetToOrignal.DotProduct(YAxis);
            FrustumPointInLightCoordinate[3].Z=OffsetToOrignal.DotProduct(ZAxis);
        }

        //left Top far
        {
            SVec3 OffsetToOrignal=Frustums[i].LeftTopFar -Orignal;
            FrustumPointInLightCoordinate[4].X=OffsetToOrignal.DotProduct(XAxis);
            FrustumPointInLightCoordinate[4].Y=OffsetToOrignal.DotProduct(YAxis);
            FrustumPointInLightCoordinate[4].Z=OffsetToOrignal.DotProduct(ZAxis);
        }

        //right top far
        {
            SVec3 OffsetToOrignal=Frustums[i].RightTopFar - Orignal;
            FrustumPointInLightCoordinate[5].X=OffsetToOrignal.DotProduct(XAxis);
            FrustumPointInLightCoordinate[5].Y=OffsetToOrignal.DotProduct(YAxis);
            FrustumPointInLightCoordinate[5].Z=OffsetToOrignal.DotProduct(ZAxis);
        }
        
        //left bottom far
        {
            SVec3 OffsetToOrignal=Frustums[i].LeftBottomFar - Orignal;
            FrustumPointInLightCoordinate[6].X=OffsetToOrignal.DotProduct(XAxis);
            FrustumPointInLightCoordinate[6].Y=OffsetToOrignal.DotProduct(YAxis);
            FrustumPointInLightCoordinate[6].Z=OffsetToOrignal.DotProduct(ZAxis);
        }

        //right bottom far
        {
            SVec3 OffsetToOrignal=Frustums[i].RightBottomFar - Orignal;
            FrustumPointInLightCoordinate[7].X=OffsetToOrignal.DotProduct(XAxis);
            FrustumPointInLightCoordinate[7].Y=OffsetToOrignal.DotProduct(YAxis);
            FrustumPointInLightCoordinate[7].Z=OffsetToOrignal.DotProduct(ZAxis);
        }

        //计算包围盒,在光源坐标系下
        SVec3 Min = FrustumPointInLightCoordinate[0];
        SVec3 Max = FrustumPointInLightCoordinate[0];

        for (uint8_t j = 1; j < 8; j++)
        {
            Min.X = std::min(Min.X, FrustumPointInLightCoordinate[j].X);
            Min.Y = std::min(Min.Y, FrustumPointInLightCoordinate[j].Y);
            Min.Z = std::min(Min.Z, FrustumPointInLightCoordinate[j].Z);

            Max.X = std::max(Max.X, FrustumPointInLightCoordinate[j].X);
            Max.Y = std::max(Max.Y, FrustumPointInLightCoordinate[j].Y);
            Max.Z = std::max(Max.Z, FrustumPointInLightCoordinate[j].Z);
        }

        SVec3 BoxSize = Max - Min;  // light space box size
        SVec3 BoxCenter = (Min + Max) / 2.0f; // light space box center

        float SizeBias= 1.f; //make sure the box is a little bigger than the frustum, to avoid gap between frustum
        SVec2 ShadowMapSizeInWorldSpace =SVec2(BoxSize.X+SizeBias,BoxSize.Y+SizeBias);

        //用包围盒反推光源位置，在光源坐标系下
        SVec3 LightPosition = BoxCenter; // light position in light space
        LightPosition.Z -= BoxSize.Z/2; //make sure the light is above the frustum
        //让光源位置再远一些，确保在视椎体之外的物体能够产生阴影
        LightPosition.Z -=  LightData.MaxCastShadowDistance; 

        //光源位置在世界坐标系下
        SVec3 LightPostionInWorldSpace =Orignal+ LightPosition.X * XAxis + LightPosition.Y * YAxis + LightPosition.Z * ZAxis;

        SMatrix4x4 ViewMatrix = SMatrix4x4::LookAt(LightPostionInWorldSpace, LightData.Rotation.GetForwardDirection(), LightData.Rotation.GetUpDirection());
        SMatrix4x4 ProjectionMatrix = SMatrix4x4::Orthographic(
            -ShadowMapSizeInWorldSpace.X / 2
            , ShadowMapSizeInWorldSpace.X / 2
            , -ShadowMapSizeInWorldSpace.Y / 2
            , ShadowMapSizeInWorldSpace.Y / 2
            , 0.1f
            ,  LightData.MaxCastShadowDistance + BoxSize.Z+1.f);
    
        ProjectViewMatrixs[i] = ProjectionMatrix * ViewMatrix;
    }

    //make sure images is created
    if(CascadeShadowMaps.Empty())
    {
        CascadeShadowMaps.Resize(ShadowMapNum);
        for (uint8_t i = 0; i < ShadowMapNum; i++)
        {
            CascadeShadowMaps[i] = std::make_shared<CGPURenderTarget>();
            CascadeShadowMaps[i]->Initialize(
                vk::Extent3D(CVulkanContext::Get()->GetSceneGlobal()->DirectionalLightShadowMapExtent,1)
                ,CVulkanContext::Get()->GetSceneGlobal()->DirectionalLightShadowMapImageFormat
                ,vk::ImageUsageFlagBits::eDepthStencilAttachment | vk::ImageUsageFlagBits::eSampled
            );
        }

        //create frame buffers
        CascadeShadowMapFrameBuffers.Resize(ShadowMapNum);
        for (uint8_t i = 0; i < ShadowMapNum; i++)
        {
            TVector<vk::ImageView> Attachments=
            {
                CascadeShadowMaps[i]->GetImageView()
            };
        

            vk::FramebufferCreateInfo FrameBufferCreateInfo = vk::FramebufferCreateInfo()
                .setRenderPass(CVulkanContext::Get()->GetSceneGlobal()->DirectionalLightShadowMapRenderPass)
                .setAttachments(Attachments)
                .setWidth(CascadeShadowMaps[i]->GetSize().width)
                .setHeight(CascadeShadowMaps[i]->GetSize().height)
                .setLayers(1);

            CascadeShadowMapFrameBuffers[i] = CVulkanContext::Get()->GetDevice().createFramebuffer(FrameBufferCreateInfo);
        }

    }

    //record command
    vk::ClearValue ClearValue = vk::ClearValue()
        .setDepthStencil(vk::ClearDepthStencilValue(1.0f, 0));

    for (auto& MeshBatch : InScene->GetMeshInstancePainter()->DirectionalLightShadowMapMeshMap[LightIndex])
    {
        MeshBatch.Value.UpdateUniformBuffer();
        MeshBatch.Value.CreateDescriptorSet();
    };

    for (uint8_t i = 0; i < ShadowMapNum; i++)
    {
        auto DirectionalLightRenderPassBeginInfo=vk::RenderPassBeginInfo()
            .setRenderPass(CVulkanContext::Get()->GetSceneGlobal()->DirectionalLightShadowMapRenderPass)
            .setClearValueCount(1)
            .setPClearValues(&ClearValue)
            .setFramebuffer(CascadeShadowMapFrameBuffers[i])
            .setRenderArea(
                vk::Rect2D(vk::Offset2D(0, 0)
                , vk::Extent2D(CascadeShadowMaps[i]->GetSize().width,CascadeShadowMaps[i]->GetSize().height)))
            ;

        //begin render pass
        InCmdBuffers.beginRenderPass(DirectionalLightRenderPassBeginInfo,vk::SubpassContents::eInline);

        //bind pipeline
        InCmdBuffers.bindPipeline(vk::PipelineBindPoint::eGraphics,CVulkanContext::Get()->GetSceneGlobal()->DirectionalLightShadowMapPipeline);

        //render all mesh instances
        for (auto& MeshBatch : InScene->GetMeshInstancePainter()->DirectionalLightShadowMapMeshMap[LightIndex])
        {
            CSceneGlobal::SDirectionalLightVertexPushConstant VertexPushConstant;
            VertexPushConstant.ProjectionViewMatrix=ProjectViewMatrixs[i];

            InCmdBuffers.pushConstants(
                CVulkanContext::Get()->GetSceneGlobal()->DirectionalLightShadowMapPipelineLayout
                ,vk::ShaderStageFlagBits::eVertex
                ,0
                ,sizeof(VertexPushConstant),&VertexPushConstant);

            for(int32_t BatchIndex=0;BatchIndex<MeshBatch.Value.ModelBufferDescriptorSets.Num();BatchIndex++)
            {
                InCmdBuffers.bindDescriptorSets(vk::PipelineBindPoint::eGraphics
                    , CVulkanContext::Get()->GetSceneGlobal()->DirectionalLightShadowMapPipelineLayout
                    , 0
                    , { MeshBatch.Value.ModelBufferDescriptorSets[BatchIndex] }
                    , {}
                );
                
                //这个小批次的大小
                int32_t BatchSize=
                    CMath::Min(MESH_BATCH_SIZE,MeshBatch.Value.ModelMatrices.Num()-BatchIndex*MESH_BATCH_SIZE);

                //绘制mesh
                MeshBatch.Value.Mesh->RecordDrawCommands(InCmdBuffers,BatchSize,0);
            }
        }
    
        //end render pass
        InCmdBuffers.endRenderPass();
    }
}

void CVulkanDirectionalLight::RecordDefferShadingCommand(
    vk::CommandBuffer InCmdBuffers
    , CVulkanScene *InScene
    , const SSceneParameters &InParameters
    , CDefferRenderAttachments *Attachments)
{

    //创建描述符
    {
        int ShadowMapNum= CVulkanContext::Get()->GetSceneGlobal()->GetCascadeNum();

        SVulkanDescriptorLayout DescriptorSetLayout= CVulkanContext::Get()->GetSceneGlobal()->DirectionalLightDefferShadingDescriptorSetLayout;


        //create descriptor set, binding 0 is shadowmaps and binding 1 is uniform buffer
        vk::DescriptorSetAllocateInfo DescriptorSetAllocateInfo = vk::DescriptorSetAllocateInfo()
            .setDescriptorPool(CVulkanContext::Get()->GetAutoResetDescriptorPool()->GenerallPool.Pool)
            .setDescriptorSetCount(1)
            .setPSetLayouts(&DescriptorSetLayout.Layout);

        DescriptorSet = CVulkanContext::Get()->GetDevice().allocateDescriptorSets(DescriptorSetAllocateInfo)[0];

        //update descriptor set
        TVector<vk::DescriptorImageInfo> DescriptorImageInfos;
        DescriptorImageInfos.Resize(ShadowMapNum);
        for (uint8_t i = 0; i < ShadowMapNum; i++)
        {
            DescriptorImageInfos[i] = vk::DescriptorImageInfo()
                .setImageLayout(vk::ImageLayout::eDepthStencilReadOnlyOptimal)
                .setImageView(CascadeShadowMaps[i]->GetImageView())
                .setSampler(CVulkanContext::Get()->GetSceneGlobal()->ShadowMapSampler)
                ;
        }

        vk::DescriptorBufferInfo DescriptorBufferInfo = vk::DescriptorBufferInfo()
            .setBuffer(DefferShadingUniformBuffer->GetBuffer())
            .setOffset(0)
            .setRange(DefferShadingUniformBuffer->GetSize());

        TVector<vk::WriteDescriptorSet> WriteDescriptorSets =
        {
            vk::WriteDescriptorSet()
            .setDstSet(DescriptorSet)
            .setDstBinding(0)
            .setDescriptorType(vk::DescriptorType::eCombinedImageSampler)
            .setDescriptorCount(ShadowMapNum)
            .setImageInfo(DescriptorImageInfos)
            
            ,
            vk::WriteDescriptorSet()
            .setDstSet(DescriptorSet)
            .setDstBinding(1)
            .setDescriptorType(vk::DescriptorType::eUniformBuffer)
            .setDescriptorCount(1)
            .setBufferInfo(DescriptorBufferInfo)
        };

        CVulkanContext::Get()->GetDevice().updateDescriptorSets(WriteDescriptorSets, {});

    }

    InCmdBuffers.bindDescriptorSets(vk::PipelineBindPoint::eGraphics
        , CVulkanContext::Get()->GetSceneGlobal()->DirectionalLightDefferShadingPipelineLayout
        , 0
        , { Attachments->DescriptorSet, DescriptorSet }
        , {}
    );

    //uniform buffer
    SDirectionalLightDefferShadingFragmentUniformBuffer UniformBuffer;
    UniformBuffer.LightColorWithIntensity = 
        LightData.Color.RGB() 
        * LightData.Intensity;
    UniformBuffer.LightDirection = LightData.Rotation.GetForwardDirection();
    UniformBuffer.ViewPos =SVec3(
        InParameters.CameraPosition[0],
        InParameters.CameraPosition[1],
        InParameters.CameraPosition[2]
    ); 
    UniformBuffer.ShadowBias =LightData.ShadowBias;

    for (int i = 0; i < EngineCascadeNum; ++i)
    {
        UniformBuffer.LightVP[i] = ProjectViewMatrixs[i];
    }
     
    for (int i = 0; i < EngineCascadeNum; ++i)
    {
        UniformBuffer.MaxDepths[i].X = CascadeMaxDepths[i];
    }
    
    DefferShadingUniformBuffer->WriteData(UniformBuffer);
    
    InCmdBuffers.draw(3, 1, 0, 0);


}
